package top.hjcwzx.wrench.limiter.strategy.algorithm.impl;

import top.hjcwzx.wrench.limiter.memory.ExpirableKeyStore;
import top.hjcwzx.wrench.limiter.strategy.algorithm.RateLimiterAlgorithmStrategy;

public class TokenBucketRateLimiterStrategy extends ExpirableKeyStore<String, TokenBucketRateLimiterStrategy.TokenBucket>
        implements RateLimiterAlgorithmStrategy {

    // 默认容量和填充速率
    private final int capacity;
    private final double refillRatePerSecond;

    public TokenBucketRateLimiterStrategy(
            EvictionPolicy evictionPolicy,
            long expireAfterAccessMs,
            int maxCacheSize,
            int capacity,
            double refillRatePerSecond) {
        super(evictionPolicy, expireAfterAccessMs, maxCacheSize);
        this.capacity = capacity;
        this.refillRatePerSecond = refillRatePerSecond;
    }

    @Override
    public boolean tryAcquire(String key, int permits) {
        TokenBucket bucket = get(key);

        if (bucket == null) {
            bucket = new TokenBucket(capacity, refillRatePerSecond);
            put(key, bucket);
        }

        return bucket.tryConsume(permits);
    }

    static class TokenBucket {
        private final int capacity;
        private final double refillRate; // 每秒恢复令牌数
        private double tokens;
        private long lastRefillTimeNanos;

        public TokenBucket(int capacity, double refillRate) {
            this.capacity = capacity;
            this.refillRate = refillRate;
            this.tokens = capacity;
            this.lastRefillTimeNanos = System.nanoTime();
        }

        public synchronized boolean tryConsume(int permits) {
            refill();

            if (tokens >= permits) {
                tokens -= permits;
                return true;
            }
            return false;
        }

        private void refill() {
            long now = System.nanoTime();
            double secondsElapsed = (now - lastRefillTimeNanos) / 1_000_000_000.0;
            double tokensToAdd = secondsElapsed * refillRate;

            if (tokensToAdd > 0) {
                tokens = Math.min(capacity, tokens + tokensToAdd);
                lastRefillTimeNanos = now;
            }
        }
    }
}
